package themis.flink.riskengine;

import com.alibaba.fastjson.JSONArray;
import com.alibaba.fastjson.JSONObject;
import java.io.IOException;
import java.io.Serializable;
import java.util.Map;
import org.apache.flink.api.common.functions.RichMapFunction;
import org.apache.flink.api.common.state.ValueState;
import org.apache.flink.api.common.state.ValueStateDescriptor;
import org.apache.flink.configuration.Configuration;

public class KeyEnrichFunction extends RichMapFunction<JSONObject, JSONObject> {
    private ValueState<Serializable> keyState;

    public void open(Configuration config) {
        this.keyState = getRuntimeContext()
                .getState(new ValueStateDescriptor("saved keyState", Serializable.class));
    }

    private <T> T getState(Class<T> tClass) throws IOException {
        return tClass.cast(this.keyState.value());
    }

    private void setState(Serializable v) throws IOException {
        this.keyState.update(v);
    }

    public JSONObject map(JSONObject event) throws Exception {
        String keyName = event.getString("KEY_NAME");
        if (keyName.equals("event"))
            return event;
        if (keyName.endsWith(".history")) {
            JSONArray history = getState(JSONArray.class);
            if (history == null)
                history = new JSONArray();
            history.add(event);
            if (history.size() > 100)
                history.remove(0);
            setState((Serializable)history);
            JSONObject newEvent = new JSONObject();
            newEvent.putAll((Map)event);
            newEvent.put("HISTORY", history);
            return newEvent;
        }
        throw new UnsupportedOperationException("unsupported key type");
    }
}
